import os
os.environ['ETS_TOOLKIT'] = 'qt4'#导入chaco时需要
import numpy as np
import scipy.special as sc
from chaco.api import *
from chaco.tools.api import DragZoom, PanTool, ZoomTool
from enable.api import ComponentEditor
from traits.api import *
from traitsui.api import *

def sigmai_mc_cormick(alpha,re):#攻角：度，雷诺数
    return (alpha*np.pi/180)**1.4*re**.35

def sigmai_mc_cormick_modified(k,cl,re):#攻角：度，雷诺数
    return k*cl**2*re**.4

def alpha_schneer_sauer(nb,rb):#rb:微米
    v=nb*4/3*np.pi*(rb/1e6)**3
    return v/(1+v)

def nb_schneer_sauer(alpha,rb):#rb:微米
    v=4/3*np.pi*(rb/1e6)**3
    return alpha/(1-alpha)/v





class BasicRelation(HasTraits):
    plot = Instance(ToolbarPlot)

    def _plot_default(self):
        plot = ToolbarPlot(ArrayPlotData())
        self.clear_plot(plot)
        return plot

    button1 = Button('test1')
    button2 = Button('test2')
    button3 = Button('test3')
    fun = Enum(
        'McCormick梢涡空化初生预报模型',
        '修正McCormick梢涡空化初生预报模型(k=.035)',
        '修正McCormick梢涡空化初生预报模型(k=.073)',
        'Schnerr-Sauer空化模型气相体积分数alpha',
        'Schnerr-Sauer空化模型气泡密度',
        '函数1',
        '函数2',
        '函数3',
    )

    def _fun_changed(self):
        plot = self.plot
        self.clear_plot(plot)
        match self.fun:
            case 'McCormick梢涡空化初生预报模型':
                x = np.linspace(1e3, 1e6, 1001)
                plot.data.update_data(
                    x=x,
                    y1=sigmai_mc_cormick(1,x),
                    y2=sigmai_mc_cormick(2,x),
                    y3=sigmai_mc_cormick(3,x),
                    y4=sigmai_mc_cormick(4,x),
                    y5=sigmai_mc_cormick(5,x),
                )
                plot.x_axis.title = "Renold Number"
                plot.y_axis.title = "Critical Cavitation Number"
                plot.plot(('x', 'y1'), name='alpha=1 degree', color='auto')
                plot.plot(('x', 'y2'), name='alpha=2 degree', color='auto')
                plot.plot(('x', 'y3'), name='alpha=3 degree', color='auto')
                plot.plot(('x', 'y4'), name='alpha=4 degree', color='auto')
                plot.plot(('x', 'y5'), name='alpha=5 degree', color='auto')
            case '修正McCormick梢涡空化初生预报模型(k=.035)':
                x = np.linspace(1e3, 1e6, 1001)
                k=.035
                plot.data.update_data(
                    x=x,
                    y1=sigmai_mc_cormick_modified(k,.4,x),
                    y2=sigmai_mc_cormick_modified(k,.6,x),
                    y3=sigmai_mc_cormick_modified(k,.8,x),
                    y4=sigmai_mc_cormick_modified(k,1,x),
                    y5=sigmai_mc_cormick_modified(k,1.2,x),
                )
                plot.x_axis.title = "Renold Number"
                plot.y_axis.title = "Critical Cavitation Number"
                plot.plot(('x', 'y1'), name='cl=.4', color='auto')
                plot.plot(('x', 'y2'), name='cl=.6', color='auto')
                plot.plot(('x', 'y3'), name='cl=.8', color='auto')
                plot.plot(('x', 'y4'), name='cl=1.0', color='auto')
                plot.plot(('x', 'y5'), name='cl=1.2', color='auto')
            case '修正McCormick梢涡空化初生预报模型(k=.073)':
                x = np.linspace(1e3, 1e6, 1001)
                k=.073
                plot.data.update_data(
                    x=x,
                    y1=sigmai_mc_cormick_modified(k,.4,x),
                    y2=sigmai_mc_cormick_modified(k,.6,x),
                    y3=sigmai_mc_cormick_modified(k,.8,x),
                    y4=sigmai_mc_cormick_modified(k,1,x),
                    y5=sigmai_mc_cormick_modified(k,1.2,x),
                )
                plot.x_axis.title = "Renold Number"
                plot.y_axis.title = "Critical Cavitation Number"
                plot.plot(('x', 'y1'), name='cl=.4', color='auto')
                plot.plot(('x', 'y2'), name='cl=.6', color='auto')
                plot.plot(('x', 'y3'), name='cl=.8', color='auto')
                plot.plot(('x', 'y4'), name='cl=1.0', color='auto')
                plot.plot(('x', 'y5'), name='cl=1.2', color='auto')
            case 'Schnerr-Sauer空化模型气相体积分数alpha':
                x = np.linspace(20, 200, 181)
                plot.data.update_data(
                    x=x,
                    y8=alpha_schneer_sauer(1e8,x),
                    y9=alpha_schneer_sauer(1e9,x),
                    y10=alpha_schneer_sauer(1e10,x),
                    y11=alpha_schneer_sauer(1e11,x),
                )
                plot.x_axis.title = "Renold Number"
                plot.y_axis.title = "vapor volume factor"
                plot.plot(('x', 'y8'), name='nb=1e8', color='auto')
                plot.plot(('x', 'y9'), name='nb=1e9', color='auto')
                plot.plot(('x', 'y10'), name='nb=1e10', color='auto')
                plot.plot(('x', 'y11'), name='nb=1e11', color='auto')
            case 'Schnerr-Sauer空化模型气泡密度':
                x = np.linspace(10, 200, 191)
                plot.data.update_data(
                    x=x,
                    y5=nb_schneer_sauer(.5,x),
                    y1=nb_schneer_sauer(.1,x),
                    y05=nb_schneer_sauer(.05,x),
                    y01=nb_schneer_sauer(.01,x),
                )
                plot.x_axis.title = "Renold Number"
                plot.y_axis.title = "Critical Cavitation Number"
                # plot.index_range.low_setting = 0
                # plot.index_range.high_setting = 200
                # plot.value_range.low_setting = 0
                # plot.value_range.high_setting = 1e12
                plot.plot(('x', 'y5'), name='alpha=0.5', color='auto')
                plot.plot(('x', 'y1'), name='alpha=0.1', color='auto')
                plot.plot(('x', 'y05'), name='alpha=0.05', color='auto')
                plot.plot(('x', 'y01'), name='alpha=0.01', color='auto')
            case '函数1':
                x = np.linspace(0, 10, 21)
                y = x
                plot.data.update_data(
                    x=x,
                    y=y,
                    y1=y + 10,
                )
                plot.x_axis.title = "x"
                plot.y_axis.title = "y"
                plot.plot(('x', 'y'), name='fjdk', color='auto')
                plot.plot(('x', 'y1'), name='fjdk1', color='auto')
            case '函数2':
                x = np.linspace(0, 20, 41)
                y = x ** 2
                plot.data.update_data(
                    x=x,
                    y=y,
                    y1=y + 10,
                )
                plot.x_axis.title = "x"
                plot.y_axis.title = "y"
                plot.plot(('x', 'y'), name='fjdk', color='auto')
                plot.plot(('x', 'y1'), name='fjdk1', color='auto')
            case '函数3':
                x = np.linspace(0, 40, 81)
                y = x ** 3
                plot.data.update_data(
                    x=x,
                    y=y,
                    y1=y + 10,
                )
                plot.x_axis.title = "x"
                plot.y_axis.title = "y"
                plot.plot(('x', 'y'), name='fjdk', color='auto')
                plot.plot(('x', 'y1'), name='fjdk1', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        plot.request_redraw()

    def traits_init(self):
        self._fun_changed()

    def clear_plot(self, plot):
        plots = [p for p in plot.plots.keys()]
        for p in plots:
            plot.delplot(p)
        for overlay in plot.overlays:
            plot.overlays.remove(overlay)
        plot.request_redraw()

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            ),
            HGroup(
                UItem('fun'),
                spring,
            ),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


class Form(HasTraits):
    basic_relation = Instance(BasicRelation, ())
    traits_view = View(
        Tabbed(
            UCustom('basic_relation', label='基本参数'),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


if __name__ == '__main__':
    Form().configure_traits()
